import es
import wcs

from UserDict import UserDict

class cRestrict(UserDict):
	def __init__(self, data = {}, **kw):
		UserDict.__init__(self)

	def restrict(self, userid, weapon):
		userid = str(userid)
		if es.exists('userid', userid):
			weapon = weaponw(weapon)
			if not userid in self.data:
				self.data[userid] = {}

			remove(userid, weapon)

			if not weapon in self.data[userid]:
				self.data[userid][weapon] = 0

			self.data[userid][weapon] += 1

	def allow(self, userid, weapon):
		userid = str(userid)
		if es.exists('userid', userid):
			weapon = weaponw(weapon)
			if not userid in self.data:
				self.data[userid] = {}

			if not weapon in self.data[userid]:
				self.data[userid][weapon] = 1

			self.data[userid][weapon] -= 1

			if not self.data[userid][weapon]:
				self.data[userid].__delitem__(weapon)
restrictData = cRestrict()

def restrict(userid, weapon):
	restrictData.restrict(userid, weapon)

def allow(userid, weapon):
	restrictData.allow(userid, weapon)

def weaponw(weapon):
	return str(weapon).lower().replace('weapon_', '')

def remove(userid, weapon):
	longname = 'weapon_'+weaponw(weapon)
	if es.createplayerlist(userid)[int(userid)]['weapon'] == longname:
		es.cexec(userid, 'lastinv')

	handle = es.getplayerhandle(userid)
	for index in es.createentitylist(longname):
		if es.getindexprop(index, 'CBaseEntity.m_hOwnerEntity') <> handle:
			continue

		saferemove(index)

		wcs.wcs.tell(userid, 'xrestrict: restricted')
		break

def saferemove(index):
	if int(index) in es.createentitylist(index):
		if not es.entitygetvalue(index, 'classname') == 'worldspawn':
			v = es.getuserid()
			if es.exists('userid', v):
				es.entitysetvalue(index, 'targetname', index)
				es.fire(v, index, 'kill')
			else:
				es.server.insertcmd('es_xremove %s'%index)
